// 在您的宏 crate 中
extern crate proc_macro;
use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, DeriveInput, FnArg, Ident, ImplItem, ItemImpl, Pat, Type};

/// 派生宏：自动实现ReflectCall trait
#[proc_macro_derive(ReflectCall)]
pub fn derive_reflect_call(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let struct_name = &input.ident;
    let struct_name_str = struct_name.to_string();

    let gen = quote! {
        impl ReflectCall for #struct_name {
            fn call_method<'a>(
                &'a mut self,
                name: &'a str,
                args: JsonValue
            ) -> BoxFuture<'a, Result<JsonValue>> {
                let reflect_methods = Self::reflected_methods();
                if let Some(method) = reflect_methods.get(name) {
                    method.invoke(self, args)
                } else {
                    Box::pin(async move {
                        Err(anyhow!("Method {} not found", name))
                    })
                }
            }

            fn as_any(&self) -> &dyn Any {
                self
            }

            fn as_any_mut(&mut self) -> &mut dyn Any {
                self
            }
        }

        impl #struct_name {
            /// 获取反射方法表 - 使用OnceCell优化性能
            pub fn reflected_methods() -> &'static HashMap<&'static str, ReflectedMethod> {
                use once_cell::sync::OnceCell;
                static INSTANCE: OnceCell<HashMap<&'static str, ReflectedMethod>> = OnceCell::new();

                INSTANCE.get_or_init(|| {
                    let mut map = HashMap::new();
                    for registration in inventory::iter::<MethodRegistration> {
                        if registration.struct_name == #struct_name_str {
                            map.insert(registration.method.name, registration.method.clone());
                        }
                    }
                    map
                })
            }
        }
    };
    gen.into()
}

/// 属性宏：自动注册反射方法
#[proc_macro_attribute]
pub fn reflect_method(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let mut item_impl = parse_macro_input!(item as ItemImpl);
    let struct_name = match &item_impl.self_ty.as_ref() {
        Type::Path(type_path) => type_path.path.segments.last().unwrap().ident.clone(),
        _ => panic!("Only implemented for simple struct paths"),
    };

    let struct_name_str = struct_name.to_string();
    let mut new_items = Vec::new();
    let mut reflection_functions = Vec::new();
    let mut registrations = Vec::new();

    // 获取原始 items 的所有权
    let items = std::mem::take(&mut item_impl.items);

    for item in items {
        if let ImplItem::Fn(method) = item {
            // 跳过静态方法（没有接收者）
            if method.sig.receiver().is_none() {
                new_items.push(ImplItem::Fn(method));
                continue;
            }


            let method_name = method.sig.ident.to_string();
            let method_ident = method.sig.ident.clone();
            let is_async = method.sig.asyncness.is_some();
            let is_mut = match method.sig.receiver() {
                Some(receiver) => receiver.mutability.is_some(),
                None => false,
            };

            // 收集参数信息
            let mut param_names = Vec::new();
            let mut param_type_tokens = Vec::new();
            let mut index = 0;

            for arg in method.sig.inputs.iter().skip(1) {
                if let FnArg::Typed(pat_type) = arg {
                    let ty = &pat_type.ty;
                    param_type_tokens.push(ty.clone());

                    if let Pat::Ident(pat_ident) = &*pat_type.pat {
                        let name = pat_ident.ident.to_string();
                        param_names.push(name);
                    } else {
                        let arg_name = format!("arg_{}", index);
                        param_names.push(arg_name);
                    }
                    index += 1;
                }
            }

            // 计算参数数量
            let param_count = param_names.len();

            // 生成参数解包代码
            let arg_unpack = if param_count == 0 {
                // 无参数
                quote! {
                    if !args.is_null() && (!args.is_object() || !args.as_object().unwrap().is_empty()) {
                        return Box::pin(async move {
                            Err(anyhow!("Expected no arguments, but got: {}", args))
                        });
                    }
                }
            } else if param_count == 1 {
                // 单个参数
                let name = &param_names[0];
                let ty = &param_type_tokens[0];
                let ident = Ident::new(name, proc_macro2::Span::call_site());

                quote! {
                    let #ident: #ty = match args.get(#name) {
                        Some(val) => match serde_json::from_value(val.clone()) {
                            Ok(v) => v,
                            Err(_) => {
                                // 尝试将整个JSON对象反序列化为参数类型
                                match serde_json::from_value(args.clone()) {
                                    Ok(v) => v,
                                    Err(e) => return Box::pin(async move {
                                        Err(anyhow!("Failed to deserialize argument '{}' or the whole object: {}", #name, e))
                                    }),
                                }
                            }
                        },
                        None => {
                            // 尝试将整个JSON对象反序列化为参数类型
                            match serde_json::from_value(args.clone()) {
                                Ok(v) => v,
                                Err(e) => return Box::pin(async move {
                                    Err(anyhow!("Failed to deserialize the whole object as argument '{}': {}", #name, e))
                                }),
                            }
                        },
                    };
                }
            } else {
                // 多个参数
                let unpack_code = param_names.iter().zip(param_type_tokens.iter()).map(|(name, ty)| {
                    let ident = Ident::new(name, proc_macro2::Span::call_site());
                    quote! {
                        let #ident: #ty = match args.get(#name) {
                            Some(val) => match serde_json::from_value(val.clone()) {
                                Ok(v) => v,
                                Err(e) => return Box::pin(async move {
                                    Err(anyhow!("Failed to deserialize argument '{}': {}", #name, e))
                                }),
                            },
                            None => return Box::pin(async move {
                                Err(anyhow!("Missing required argument: {}", #name))
                            }),
                        };
                    }
                });
                quote! { #(#unpack_code)* }
            };

            // 生成方法调用表达式
            let method_call = if param_count > 0 {
                let idents = param_names
                    .iter()
                    .map(|name| Ident::new(name, proc_macro2::Span::call_site()));
                quote! { this.#method_ident(#(#idents),*) }
            } else {
                quote! { this.#method_ident() }
            };

            // 判断返回类型
            let return_type = &method.sig.output;
            let is_result = match return_type {
                syn::ReturnType::Type(_, ty) => {
                    if let syn::Type::Path(type_path) = &**ty {
                        type_path.path.segments.last().map_or(false, |seg| seg.ident == "Result")
                    } else {
                        false
                    }
                }
                _ => false,
            };
            let is_unit = matches!(return_type, syn::ReturnType::Default);

            // 为每个方法生成一个独立的函数
            let function_name = Ident::new(
                &format!("__reflect_{}_{}", struct_name, method_name),
                proc_macro2::Span::call_site(),
            );

            // 生成结果处理代码
            let function_impl = if is_async {
                // 异步方法处理
                let result_handling = if is_result {
                    quote! {
                        match result.await {
                            Ok(val) => serde_json::to_value(&val)
                                .map_err(|e| anyhow!("Result serialization failed: {}", e)),
                            Err(err) => Err(anyhow!("Method error: {}", err)),
                        }
                    }
                } else if is_unit {
                    quote! {
                        result.await; // 执行但不使用结果
                        Ok(serde_json::Value::Null)
                    }
                } else {
                    quote! {
                        let val = result.await;
                        serde_json::to_value(&val)
                            .map_err(|e| anyhow!("Return value serialization failed: {}", e))
                    }
                };

                quote! {
                    // 在函数内部定义 param_count
                    let param_count = #param_count;

                    // 对于无参数方法，不需要检查是否为对象
                    if param_count > 0 && !args.is_object() {
                        return Box::pin(async move {
                            Err(anyhow!("Expected JSON object for arguments"))
                        });
                    }

                    #arg_unpack

                    let result = #method_call;
                    Box::pin(async move {
                        #result_handling
                    })
                }
            } else {
                // 同步方法处理
                let result_handling = if is_result {
                    quote! {
                        match result {
                            Ok(val) => serde_json::to_value(&val)
                                .map_err(|e| anyhow!("Result serialization failed: {}", e)),
                            Err(err) => Err(anyhow!("Method error: {}", err)),
                        }
                    }
                } else if is_unit {
                    quote! {
                        result; // 执行但不使用结果
                        Ok(serde_json::Value::Null)
                    }
                } else {
                    quote! {
                        serde_json::to_value(&result)
                            .map_err(|e| anyhow!("Return value serialization failed: {}", e))
                    }
                };

                quote! {
                    // 在函数内部定义 param_count
                    let param_count = #param_count;

                    // 对于无参数方法，不需要检查是否为对象
                    if param_count > 0 && !args.is_object() {
                        return Box::pin(async move {
                            Err(anyhow!("Expected JSON object for arguments"))
                        });
                    }

                    #arg_unpack

                    let result = #method_call;
                    Box::pin(async move {
                        #result_handling
                    })
                }
            };

            // 生成完整的函数定义
            let function_def = if is_mut {
                // 可变方法版本
                quote! {
                    fn #function_name<'a>(
                        target: &'a mut dyn ReflectCall,
                        args: JsonValue
                    ) -> BoxFuture<'a, Result<JsonValue>> {
                        let this = match target.as_any_mut().downcast_mut::<#struct_name>() {
                            Some(t) => t,
                            None => return Box::pin(async move {
                                Err(anyhow!("Downcast failed"))
                            }),
                        };
                        #function_impl
                    }
                }
            } else {
                // 不可变方法版本
                quote! {
                    fn #function_name<'a>(
                        target: &'a mut dyn ReflectCall,
                        args: JsonValue
                    ) -> BoxFuture<'a, Result<JsonValue>> {
                        // 注意：这里使用不可变引用
                        let this = match target.as_any().downcast_ref::<#struct_name>() {
                            Some(t) => t,
                            None => return Box::pin(async move {
                                Err(anyhow!("Downcast failed"))
                            }),
                        };
                        #function_impl
                    }
                }
            };

            // 生成方法注册代码
            let registration = quote! {
                inventory::submit! {
                    MethodRegistration {
                        struct_name: #struct_name_str,
                        method: ReflectedMethod {
                            name: #method_name,
                            param_count: #param_count,
                            is_async: #is_async,
                            function: #function_name,
                        }
                    }
                }
            };

            // 保留原始方法定义
            new_items.push(ImplItem::Fn(method));

            // 添加函数定义和注册代码
            reflection_functions.push(function_def);
            registrations.push(registration);
        } else {
            // 保留其他 impl 项
            new_items.push(item);
        }
    }

    // 将新 items 放回 item_impl
    item_impl.items = new_items;

    // 组合所有生成的代码
    let output = quote! {
        #item_impl

        #(#reflection_functions)*

        #(#registrations)*
    };

    output.into()
}